import time
from os import times

import numpy as np
import pandas as pd
from tensorflow.python.ops.numpy_ops.np_random import random

from bins import vehicleParams
from bins.Map import Map
from bins.PlannerLogger import PlannerLogger
from src.motionplanner.Constrained_ILQR.scripts.ilqr.iLQR import iLQR

Logger = PlannerLogger()

class CarPlanner:
    def __init__(self, args, name):
        self.args = args
        self.name = name
        self.map = Map()
        self.vehicle = vehicleParams.VehicleParams()
        self.npc_trajectory = {}
        self.state = None
        self.current_state = None
        self.historical_track = []
        self.restrain_car = None
        self.ilqr_planner = iLQR(self.args, self.vehicle.get_BOX())
        self.desired_y = None
        self.v = None
        self.ref_line = None
        self.time_consuming = 0.0

    def set_begin(self, x, y, v, h):
        self.current_state = np.array([x, y, v, h])
        self.historical_track.append(self.current_state)

    def load_trajectory(self, filename: str = None, trajectory: np.array = None):
        if filename is not None:
            df = pd.read_csv(filename)
            npc_data = df.to_numpy()
            self.state = npc_data.T[:, 0:self.args.horizon]
        if trajectory is not None:
            self.state = trajectory

    def update_current_state(self, index = 1):
        state = self.state[index]
        state[2] = state[2] + (random() - 0.5) / 30
        self.current_state =state
        self.historical_track.append(self.state[index])

    def update_refline(self, desired_y, v = 0.0):
        self.ref_line = []
        self.desired_y = desired_y
        self.v = v
        begin_x = 0
        for i in range(0, self.map.road_length):
            self.ref_line.append(np.array([begin_x, desired_y, v]))
            begin_x += v * self.args.timestep
        self.ref_line = np.array(self.ref_line)
        self.ilqr_planner.set_global_plan(self.ref_line)

    def update_reference(self, trajectory, v, desired_y):
        self.ref_line = []
        last_index = None
        for i in trajectory:
            self.ref_line.append(np.array([i[0], i[1], v]))
            last_index = i
        begin_x = last_index[0]
        for i in range(0, self.map.road_length):
            begin_x += v * self.args.timestep
            self.ref_line.append(np.array([begin_x, desired_y, v]))

        self.ref_line = np.array(self.ref_line)
        self.ilqr_planner.set_global_plan(self.ref_line)



    def update_npc_trajectory(self, trajectory):
        self.npc_trajectory = trajectory

    def get_ego_states(self):
        ego_states = np.array([[self.current_state[0], self.current_state[1],                         0],
                               [self.current_state[2],                         0,                         0],
                               [                        0,                         0, self.current_state[3]],
                               [                        0,                         0,                         0],
                               [                        0,                         0,                         0]])
        return ego_states

    def load_restrained_vehicle(self, name, trajectory: np.array = None, distance = 10.0):
        self.args.offset = distance
        if trajectory is None:
            self.npc_trajectory[name] = trajectory
        self.restrain_car = name
        self.ilqr_planner.target_traj = self.npc_trajectory[name]

    def update_control(self):
        if self.restrain_car is not None and self.npc_trajectory[self.restrain_car] is not None:
            self.ilqr_planner.target_traj = self.npc_trajectory[self.restrain_car]
        begin_time = time.time()
        desired_path, local_plan, control, X = self.ilqr_planner.run_step(self.get_ego_states(),
                                                                           list(self.npc_trajectory.values()))
        end_time = time.time()
        self.time_consuming = end_time - begin_time
        self.state = X
        self.record_data()

    def record_data(self):
        state = {"x": list(self.state[:, 0]),
                 "y": list(self.state[:, 1]),
                 "v": list(self.state[:, 2]),
                 "h": list(self.state[:, 3])}
        cula_data = {"goal_v": self.v,
                "goal_car": self.restrain_car,
                "time_cost": self.time_consuming,
                "state": state}

        current_state = {
            "x": self.current_state[0],
            "y": self.current_state[1],
            "v": self.current_state[2],
            "h": self.current_state[3],
        }
        track_data = {
            "current_state": current_state
        }
        Logger.update_data(self.name, state=cula_data, track=track_data)
        Logger.update_args(self.name, self.args)


